#!/usr/bin/python3

# Copyright (c) 2020-2021 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import procfs
import pathlib


class Snapshot:

    def __init__(self):
        self._ps = None
        self._badPids = set()
        self._usedPids = set()
        self._usedEnvVars = dict()

        self._ps = procfs.pidstats()
        for pid in self._ps.processes.keys():
            v = self._ps.processes[pid]

            # ignore kernel threads
            if "PF_KTHREAD" in v["stat"].process_flags():
                del self._ps.processes[pid]
                continue

            # ignore user process that are started by kernel (not started by init), for example bpfilter_umh
            if pid != 1 and 1 not in self.get_anscetor_pids(pid):
                del self._ps.processes[pid]
                continue

            # load all sub objects for v (v is procfs.process object)
            try:
                v["stat"]
                v["status"]
                v["cmdline"]
                v["threads"]        # FIXME：perhaps we don't need it
                v["cgroups"]
                # procfs module would raise exception when garbage is in /proc/X/environ (maybe caused by memory corruption)
                # we do this check to prevent it
                if b'=' not in pathlib.Path("/proc/%d/environ" % (pid)).read_bytes():
                    self._badPids.add(pid)
                    continue
                v["environ"]
            except FileNotFoundError:
                self._badPids.add(pid)
                continue

    @property
    def pids(self):
        return self._procDict.keys()

    def get_process_cmdline(self, pid):
        return self._ps.processes[pid]["cmdline"]

    def get_process_environ(self, pid):
        return self._ps.processes[pid]["environ"]

    def get_process_eff_pid(self, pid):
        v = self._ps.processes[pid]
        eff_pid = pid
        if pid != v["status"]["Tgid"]:
            eff_pid = v["status"]["Tgid"]
        return eff_pid

    def get_anscetor_pids(self, pid):
        ret = []
        while True:
            pid = self._ps.processes[pid]["stat"]["ppid"]
            if pid == 0:
                return ret
            ret.append(pid)

    def get_descendant_pids(self, pid):
        ret = []
        for x in self._ps.processes.keys():
            if pid in self.get_anscetor_pids(x):
                ret.append(x)
        return ret

    def query_pids(self, filter_functions=[]):
        ret = []
        for pid in self._ps.processes.keys():
            if not self._pidUsed(pid) and not _filteredOut(pid, filter_functions):
                ret.append(pid)
        return ret

    def use_pids(self, pids, strict=True):
        for pid in pids:
            if strict:
                assert not self._pidUsed(pid)
            self._usePid(pid)

    def use_pid(self, pid):
        assert not self._pidUsed(pid)
        self._usePid(pid)

    def remove_pids(self, pids):
        for pid in pids:
            self.remove_pid(pid)

    def remove_pid(self, pid):
        # assert must remove leaf node
        assert pid not in self._usedPids
        assert pid not in self._usedEnvVars

        del self._ps.processes[pid]
        if pid in self._badPids:
            self._badPids.remove(pid)

    def has_env_var(self, pid, varName):
        return varName in self.get_process_environ(pid)

    def check_env_var(self, pid, varName, varValue):
        if varName not in self.get_process_environ(pid):
            return False
        if self.get_process_environ(pid)[varName] != varValue:
            return False
        return True

    def use_env_vars(self, pid, varNames, strict=True):
        for varName in varNames:
            if strict:
                assert not self._envVarUsed(varName)
            self._useEnvVar(varName)

    def use_env_var(self, pid, varName):
        assert not self._envVarUsed(pid, varName)

    def _pidUsed(self, pid):
        return pid in self._usedPids

    def _usePid(self, pid):
        assert pid in self._ps.processes
        assert pid not in self._usedPids
        self._usedPids[pid] = True

    def _envVarUsed(self, pid, varName):
        if pid not in self._usedEnvVars:
            return False
        if varName not in self._usedEnvVars[pid]:
            return False
        return True

    def _useEnvVar(self, pid, varName):
        assert pid in self._ps.processes and varName in self._ps.processes[pid]["environ"]
        self._usedEnvVars.put(pid, {})[varName] = True


def _filteredOut(item, filter_functions):
    for func in filter_functions:
        if not func(item):
            return True
    return False

